#include "hypercube_hmm_engine.hpp"
#include "hypercube_hmm_maxk.hpp"
#include "hypercube_hmm_midk.hpp"
#include "hypercube_hmm_mink.hpp"
#include "hypercube_hmmsq.hpp"

static constexpr hmm_t schmms[UNDEF_MATRIX_TYPE] = {hhmm_sq, hhmm_maxk,
                                                    hhmm_mink, hhmm_midk};
static constexpr hmm_t bsgs_schmms[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_bsgs, hhmm_maxk_bsgs, hhmm_mink_bsgs, hhmm_midk_bsgs};
static constexpr hmm_mt_t mt_schmms[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_mt, hhmm_maxk_mt, hhmm_mink_mt, hhmm_midk_mt};
static constexpr hmm_mt_t mt_bsgs_schmms[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_bsgs_mt, hhmm_maxk_bsgs_mt, hhmm_mink_bsgs_mt, hhmm_midk_bsgs_mt};
using st_hmm_table_t = decltype(&schmms);
using mt_hmm_table_t = decltype(&mt_schmms);
static constexpr st_hmm_table_t st_hmm_table[METHOD::METHOD_NUM] = {
    &schmms, &bsgs_schmms};
static constexpr mt_hmm_table_t mt_hmm_table[METHOD::METHOD_NUM] = {
    &mt_schmms, &mt_bsgs_schmms};

static constexpr hmm_ra_t ra_schmms_A[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_rotate_align_A, nullptr, hhmm_mink_rotate_align_A};
static constexpr hmm_ra_t ra_schmms_B[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_rotate_align_B, nullptr, hhmm_mink_rotate_align_B};

using hhmm_ra_batch_t = void (*)(hmm_constant_multipler *,
                                 const helib::EncryptedArray &, int, int,
                                 helib::Ctxt &, const helib::Ctxt &);
static constexpr hhmm_ra_batch_t ra_batch_A[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_rotate_align_A_batch};
static constexpr hhmm_ra_batch_t ra_batch_B[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_rotate_align_B_batch};

using hhmm_sc_batch_t = void (*)(hmm_constant_multipler *,
                                 const helib::EncryptedArray &, int,
                                 std::vector<helib::Ctxt> &,
                                 const helib::Ctxt &);
static constexpr hhmm_sc_batch_t sc_batch_A[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_shift_compute_A_batch};
static constexpr hhmm_sc_batch_t sc_batch_B[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_shift_compute_B_batch};

using hhmm_sa_t = helib::Ctxt (*)(
    hmm_constant_multipler *, const helib::EncryptedArray &,
    const std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> &, int);
static constexpr hhmm_sa_t sc_schmms_A[UNDEF_MATRIX_TYPE] = {
    hhmm_sq_shift_compute_A, nullptr, hhmm_mink_shift_compute_A};

static hmm_constant_multipler *
create_constant_multipler(hmm_metric_collector *collector, MATRIX_TYPE type,
                          const helib::EncryptedArray &ea,
                          const cube_signature &cube) {
  ASSERT_PRINTF(collector != nullptr, "Invalid metric_collector_\n");

  switch (type) {
  case MATRIX_TYPE::SQUARE:
    return new hhmmsq_constant_multipler(ea, collector);
  case MATRIX_TYPE::MAX_K:
    return build_hhmm_maxk_constant_multipler(collector, ea, cube);
  case MATRIX_TYPE::MIN_K:
    return build_hhmm_mink_constant_multipler(collector, ea, cube);
  case MATRIX_TYPE::MEDIUM_K:
    return build_hhmm_midk_constant_multipler(collector, ea, cube);
  default:
    ERROR_PRINTF("UNDEFINED MATRIX_TYPE: %d\n", type);
    return nullptr;
  }
}

hmm_status<hypercube_hmme>::hmm_status(const hmm_status<shmm_engine> &status)
    : hmm_status<shmm_engine>(status.method(), status.threads(),
                              status.reset()) {}

hmm_status<hypercube_hmme>::operator hmm_t() const {
  return (*st_hmm_table[method_])[type_];
}

hmm_status<hypercube_hmme>::operator hmm_mt_t() const {
  return (*mt_hmm_table[method_])[type_];
}

hypercube_hmme::hypercube_hmme(const hmm_status<shmm_engine> &status)
    : shmm_engine(new hmm_status<hypercube_hmme>(status)) {
#ifdef OMP_CTXT_REDUCTION
  INFO_PRINTF("OMP_CTXT_REDUCTION is enabled\n");
#endif
}

hypercube_hmme::hypercube_hmme(const hypercube_hmme &hmme)
    : shmm_engine(new hmm_status<hypercube_hmme>(*hmme.status())) {
  for (int i = 0; i < MATRIX_TYPE::UNDEF_MATRIX_TYPE; i++) {
    metric_collector_[i] = create_metric(static_cast<MATRIX_TYPE>(i));
  }
  // TODO: copy cunstrutor/method for `constant_multipler`
  constant_multipler_ = hmme.constant_multipler_;
}

cube_signature hypercube_hmme::get_hypercube(long m, long n, long k) {
  auto hyper_mnk = get_expected_mnk(m, n, k);
  long hyper_row = std::max(hyper_mnk.size(0), hyper_mnk.size(2));
  long hyper_col = std::max(hyper_mnk.size(2), hyper_mnk.size(1));
  return std::move(cube_signature({hyper_row, hyper_col}));
}

void hypercube_hmme::register_engine(const helib::EncryptedArray &ea,
                                     const cube_signature &cube) {
  auto type = parse_matrix_type(cube);
  ea_[type] = &ea;
  metric_collector_[type] = create_metric(type);
  metric_collector_[type]->start_hmm_encoding(); //! Encoding Stage - Start
  constant_multipler_[type] =
      create_constant_multipler(metric_collector_[type], type, ea, cube);
  metric_collector_[type]->stop_hmm_encoding(); //! Encoding Stage - Stop
}

void hypercube_hmme::cgemm(MATRIX_TYPE type, helib::Ctxt &AB,
                           const helib::Ctxt &A, const helib::Ctxt &B) {
  // sanity check
  status_->matrix_type(type);
  ASSERT_PRINTF(type < MATRIX_TYPE::UNDEF_MATRIX_TYPE,
                "Invalid matrix type = %d\n", type);
  if (metric_collector_[type] == nullptr ||
      constant_multipler_[type] == nullptr) {
    INFO_PRINTF("Must register hmm engine for %s first\n",
                MATRIX_TYPE_STR[type]);
  }

  if (status_->reset()) {
    metric_collector_[type]->reset_computing(); //! reset first
  }

  constant_multipler_[type]->set_no_encoding_reuse(!status()->encoding_reuse());

  if (type != MATRIX_TYPE::SQUARE) {
    switch (type) {
    case MATRIX_TYPE::MAX_K:
      if (status()->clean_left()) {
        set_hhmm_maxk_constant_multipler_clean_left(constant_multipler_[type]);
        status()->set_clean_left(false);
      }
      if (status()->clean_right()) {
        set_hhmm_maxk_constant_multipler_clean_right(constant_multipler_[type]);
        status()->set_clean_right(false);
      }
      break;
    case MATRIX_TYPE::MIN_K:
      if (status()->clean_left()) {
        set_hhmm_mink_constant_multipler_clean_left(constant_multipler_[type]);
        status()->set_clean_left(false);
      }
      if (status()->clean_right()) {
        set_hhmm_mink_constant_multipler_clean_right(constant_multipler_[type]);
        status()->set_clean_right(false);
      }
      break;
    case MATRIX_TYPE::MEDIUM_K:
      if (status()->clean_left()) {
        set_hhmm_midk_constant_multipler_clean_left(constant_multipler_[type]);
        status()->set_clean_left(false);
      }
      if (status()->clean_right()) {
        set_hhmm_midk_constant_multipler_clean_right(constant_multipler_[type]);
        status()->set_clean_right(false);
      }
      break;
    default:
      ERROR_PRINTF("UNDEFINED MATRIX_TYPE: %d\n", type);
    }
  }

  metric_collector_[type]->start_hmm_computing(); //! Computing Stage - Start
  if (status_->threads() > 1 || status_->threads() == 0) {
    (*status())(metric_collector_[type], constant_multipler_[type],
                status_->threads(), *ea_[type], AB, A, B);
  } else {
    (*status())(metric_collector_[type], constant_multipler_[type], *ea_[type],
                AB, A, B);
  }
  metric_collector_[type]->stop_hmm_computing(); //! Computing Stage - Stop
}

void hypercube_hmme::rotate_align_A(MATRIX_TYPE type, int rot_pos_A,
                                    int rot_size_A, helib::Ctxt &Ap,
                                    const helib::Ctxt &A) {
  ra_batch_A[type](constant_multipler_[type], *ea_[type], rot_pos_A, rot_size_A,
                   Ap, A);
}
void hypercube_hmme::rotate_align_B(MATRIX_TYPE type, int rot_pos_B,
                                    int rot_size_B, helib::Ctxt &Bp,
                                    const helib::Ctxt &B) {
  ra_batch_B[type](constant_multipler_[type], *ea_[type], rot_pos_B, rot_size_B,
                   Bp, B);
}
void hypercube_hmme::shift_compute_A(MATRIX_TYPE type, int rot_pos_A,
                                     std::vector<helib::Ctxt> &Ais,
                                     const helib::Ctxt &A) {
  sc_batch_A[type](constant_multipler_[type], *ea_[type], rot_pos_A, Ais, A);
}
void hypercube_hmme::shift_compute_B(MATRIX_TYPE type, int rot_pos_B,
                                     std::vector<helib::Ctxt> &Bis,
                                     const helib::Ctxt &B) {
  sc_batch_B[type](constant_multipler_[type], *ea_[type], rot_pos_B, Bis, B);
}

void hypercube_hmme::rotate_align_A(MATRIX_TYPE type, helib::Ctxt &Ap,
                                    const helib::Ctxt &A) {
  ra_schmms_A[type](constant_multipler_[type], *ea_[type], Ap, A);
}

void hypercube_hmme::rotate_align_B(MATRIX_TYPE type, helib::Ctxt &Bp,
                                    const helib::Ctxt &B) {
  ra_schmms_B[type](constant_multipler_[type], *ea_[type], Bp, B);
}

helib::Ctxt hypercube_hmme::shift_compute_A(
    MATRIX_TYPE type,
    const std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> &hoisted_A,
    int r) {
  return sc_schmms_A[type](constant_multipler_[type], *ea_[type], hoisted_A, r);
}